{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2f9a797f",
   "metadata": {},
   "source": [
    "# 合并和分割 Merge or split\n",
    "\n",
    "- Cat\n",
    "\n",
    "- Stack\n",
    "\n",
    "- Split\n",
    "\n",
    "- Chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2750bf66",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edaa1405",
   "metadata": {},
   "source": [
    "### concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b15b7c6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([9, 32, 8])\n"
     ]
    }
   ],
   "source": [
    "a = torch.rand(4, 32, 8)\n",
    "b = torch.rand(5, 32, 8)\n",
    "\n",
    "# dim=0  0维度索引上合并\n",
    "print(torch.cat([a, b], dim=0).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1b1a5aff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([9, 3, 32, 32])\n",
      "torch.Size([4, 4, 32, 32])\n",
      "torch.Size([4, 3, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "a1 = torch.rand(4, 3, 32, 32)\n",
    "a2 = torch.rand(5, 3, 32, 32)\n",
    "\n",
    "print(torch.cat([a1, a2], dim=0).shape)\n",
    "\n",
    "a2 = torch.rand(4, 1, 32, 32)\n",
    "# 除了0维度以外，其他维度必须配对\n",
    "# print(torch.cat([a1, a2], dim=0).shape)\n",
    "# RuntimeError: invalid argument 0:\n",
    "# Size of tensors must match except in dimension 0\n",
    "\n",
    "print(torch.cat([a1, a2], dim=1).shape)\n",
    "\n",
    "a1 = torch.rand(4, 3, 16, 32)\n",
    "a2 = torch.rand(4, 3, 16, 32)\n",
    "\n",
    "print(torch.cat([a1, a2], dim=2).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "092ac4f5",
   "metadata": {},
   "source": [
    "### stack 合并\n",
    "\n",
    "stack会创建一个新的维度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1e0a049e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 3, 32, 32])\n",
      "torch.Size([4, 3, 2, 16, 32])\n",
      "torch.Size([2, 32, 8])\n"
     ]
    }
   ],
   "source": [
    "print(torch.cat([a1, a2], dim=2).shape)\n",
    "\n",
    "# 在维度索引为2的前面插入新的维度，区分合并的两个维度\n",
    "print(torch.stack([a1, a2], dim=2).shape)\n",
    "\n",
    "a = torch.rand(32, 8)\n",
    "b = torch.rand(32, 8)\n",
    "\n",
    "print(torch.stack([a, b], dim=0).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfcf0390",
   "metadata": {},
   "source": [
    "### split 拆分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "62d85d4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 8])\n",
      "torch.Size([32, 8])\n",
      "torch.Size([2, 32, 8])\n",
      "torch.Size([1, 32, 8])\n",
      "torch.Size([1, 32, 8])\n",
      "torch.Size([1, 32, 8])\n",
      "torch.Size([1, 32, 8])\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "not enough values to unpack (expected 2, got 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[8], line 18\u001b[0m\n\u001b[0;32m     15\u001b[0m \u001b[38;5;28mprint\u001b[39m(bb\u001b[38;5;241m.\u001b[39mshape)\n\u001b[0;32m     17\u001b[0m \u001b[38;5;66;03m# 只能拆成1个，所以返回1个tensor，不能用2个tensor接受\u001b[39;00m\n\u001b[1;32m---> 18\u001b[0m aa, bb \u001b[38;5;241m=\u001b[39m c\u001b[38;5;241m.\u001b[39msplit(\u001b[38;5;241m2\u001b[39m, dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m     19\u001b[0m \u001b[38;5;66;03m# ValueError: not enough values to unpack (expected 2, got 1)\u001b[39;00m\n",
      "\u001b[1;31mValueError\u001b[0m: not enough values to unpack (expected 2, got 1)"
     ]
    }
   ],
   "source": [
    "b = torch.rand(32, 8)\n",
    "print(a.shape)\n",
    "print(b.shape)\n",
    "\n",
    "c = torch.stack([a, b], dim=0)\n",
    "print(c.shape)\n",
    "\n",
    "# split第一个参数为拆分后的维度长度\n",
    "aa, bb = c.split([1, 1], dim=0)\n",
    "print(aa.shape)\n",
    "print(bb.shape)\n",
    "\n",
    "aa, bb = c.split(1, dim=0)\n",
    "print(aa.shape)\n",
    "print(bb.shape)\n",
    "\n",
    "# 只能拆成1个，所以返回1个tensor，不能用2个tensor接受\n",
    "aa, bb = c.split(2, dim=0)\n",
    "# ValueError: not enough values to unpack (expected 2, got 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "098ded8c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
